/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements...
 */

package org.apache.flink.connector.elasticsearch.sink;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;

import java.io.ByteArrayOutputStream;
import java.io.DataInputStream;
import java.io.DataOutputStream;

/**
 * Serializes/deserializes {@link Operation} with Kryo.
 * Key fixes:
 * - Bind Kryo to the user-code/plugin ClassLoader so co.elastic.* is always visible at restore time.
 * - Use writeClassAndObject/readClassAndObject to carry runtime type information (BulkOperation variants).
 */
public class OperationSerializer {

    private final Kryo kryo;
//    private final ClassLoader userCodeCl;

    /** Default: bind to TCCL of the current thread (Task thread in runtime). */
    public OperationSerializer() {
        this(Thread.currentThread().getContextClassLoader());
    }

    /** Explicit CL injection if you prefer passing Environment#getUserCodeClassLoader(). */
    public OperationSerializer(ClassLoader cl) {
        this.kryo = new com.esotericsoftware.kryo.Kryo();

        ClassLoader ucl = (cl != null) ? cl : OperationSerializer.class.getClassLoader();
        this.kryo.setClassLoader(ucl);
        this.kryo.setRegistrationRequired(false);

        // 使用 Kryo 5 的 DefaultInstantiatorStrategy + Objenesis 的 StdInstantiatorStrategy
        // —— 全限定名，避免 “无法解析符号” 的编译错误
        //草泥马的Flink
        this.kryo.setInstantiatorStrategy(
                new com.esotericsoftware.kryo.util.DefaultInstantiatorStrategy(
                        new org.objenesis.strategy.StdInstantiatorStrategy()
                )
        );
    }

    public void serialize(Operation request, DataOutputStream out) {
        try (Output output = new Output(out)) {
            // Carry runtime type info to handle BulkOperation variants safely.
            kryo.writeClassAndObject(output, request);
            output.flush();
        }
    }

    public Operation deserialize(long requestSize, DataInputStream in) {
        try (Input input = new Input(in, (int) requestSize)) {
            return (Operation) kryo.readClassAndObject(input);
        }
    }

    public int size(Operation operation) {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        try (Output output = new Output(baos)) {
            kryo.writeClassAndObject(output, operation);
            output.flush();
            return (int) output.total();
        }
    }
}
